package mosdi.paa;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.Map;
import java.util.Queue;
import java.util.TreeMap;

import mosdi.fa.AutomatonBasedTextModel;

public abstract class TextBasedPAA {
	
	private TextBasedPAA outerThis;
	private InnerPAA paa;
	private AutomatonBasedTextModel textModel;
	private int alphabetSize;
	private ProductState[] productStates;
	private ArrayList<int[]> targets;
	private ArrayList<double[]> targetProbabilities;

	// abstract methods to be implemented by derived classes 
	public abstract double emissionProbability(int state, int emission);
	public abstract int getEmissionCount();
	public abstract int getStartState();
	public abstract int getStartValue();
	public abstract int getStateCount(); 
	public abstract int getValueCount(); 
	public abstract int performOperation(int state, int value, int emission); 
	public abstract int transitionTarget(int state, int character);

	/** Inner class implementing the underlying PAA. */
	private class InnerPAA extends PAA {
		@Override
		public int getStateCount() {
			return productStates.length;
		}
		@Override
		public double transitionProbability(int state, int targetState) {
			int[] t = targets.get(state);
			for (int i=0; i<t.length; ++i) {
				if (t[i]==targetState) return targetProbabilities.get(state)[i];
			}
			return 0.0;
		}
		@Override
		public int performOperation(int state, int value, int emission) {
			return outerThis.performOperation(productStates[state].outerState, value, emission);
		}
		@Override
		protected double[] getTargetProbabilities(int state) {
			return targetProbabilities.get(state);
		}
		@Override
		protected int[] getTargets(int state) {
			return targets.get(state);
		}
		@Override
		public double[][] stateValueStartDistribution() {
			double[][] result = new double[getStateCount()][getValueCount()];
			for (int i=0; i<textModel.getStateCount(); ++i) {
				result[i][getStartValue()] = textModel.getEquilibriumProbability(i); 
			}
			return result;
		}
		@Override
		public double emissionProbability(int state, int emission) {
			return outerThis.emissionProbability(productStates[state].outerState, emission);
		}
		@Override
		public int getEmissionCount() { return outerThis.getEmissionCount(); }
		@Override
		public int getStartState() { return outerThis.getStartState(); }
		@Override
		public int getStartValue() { return outerThis.getStartValue(); }
		@Override
		public int getValueCount() { return outerThis.getValueCount(); }
	}

	private static class ProductState {
		int outerState;
		int textModelState;
		ProductState(int outerState, int textModelState) {
			this.outerState = outerState;
			this.textModelState = textModelState;
		}
		@Override
		public boolean equals(Object obj) {
			if (!(obj instanceof ProductState)) return false;
			ProductState p = (ProductState)obj; 
			return (p.outerState==outerState) && (p.textModelState==textModelState);
		}
		@Override
		public int hashCode() {
			return textModelState ^ Integer.reverse(outerState);
		}
		@Override
		public String toString() {
			return "("+outerState+","+textModelState+")";
		}
	}
	
	public TextBasedPAA(AutomatonBasedTextModel textModel) {
		outerThis = this;
		this.textModel = textModel;
		this.alphabetSize = textModel.getAlphabetSize();
	}

	private void init() {
		buildProductAutomaton();
		paa = new InnerPAA();
	}
	
	/** Checks which product states are reachable and computes arrays
	 *  productStates, targets, and targetProbabilities. */
	private void buildProductAutomaton() {
		// maps product state on their new index.
		Map<ProductState,Integer> states = new HashMap<ProductState,Integer>();
		Queue <ProductState> queue = new LinkedList<ProductState>();
		targets = new ArrayList<int[]>();
		targetProbabilities = new ArrayList<double[]>();
		int n = 0;
		for (int i=0; i<textModel.getStateCount(); ++i) {
			ProductState newState = new ProductState(getStartState(),i);
			states.put(newState, n++);
			queue.add(newState);
		}
		while (!queue.isEmpty()) {
			ProductState state = queue.remove();
			// maps target states to probabilities
			Map<Integer,Double> transitionMap = new TreeMap<Integer,Double>();
			for (int c=0; c<alphabetSize; ++c) {
				int outerTargetState = transitionTarget(state.outerState, c);
				int textModelTargetState = textModel.getTransitionTarget(state.textModelState, c);
				ProductState targetState = new ProductState(outerTargetState, textModelTargetState);
				int targetStateIndex;
				if (states.containsKey(targetState)) {
					targetStateIndex = states.get(targetState);
				} else {
					targetStateIndex = n++;
					states.put(targetState,targetStateIndex);
					queue.add(targetState);
				}
				double p = 0.0;
				if (transitionMap.containsKey(targetStateIndex)) {
					p = transitionMap.get(targetStateIndex);
				}
				transitionMap.put(targetStateIndex, p+textModel.getProbability(state.textModelState, c));
			}
			int[] targetArray = new int[transitionMap.size()];
			double[] targetProbabilitiesArray = new double[transitionMap.size()];
			int i = 0;
			for (Map.Entry<Integer,Double> e : transitionMap.entrySet()) {
				targetArray[i] = e.getKey();
				targetProbabilitiesArray[i] = e.getValue();
				i+=1;
			}
			targets.add(targetArray);
			targetProbabilities.add(targetProbabilitiesArray);
		}
		productStates = new ProductState[states.size()];
		for (Map.Entry<ProductState,Integer> e : states.entrySet()) {
			productStates[e.getValue()] = e.getKey();
		}
	}
	
	public int getInnerStateCount() {
		return productStates.length;
	}
	
	public double[] computeValueDistribution(int iterations) {
		if (paa==null) init();
		return paa.computeValueDistribution(iterations);
	}

	/** Computes the value distribution resulting from one fixed text, i.e. the 
	 *  text model is not used here.
	 */
	public double[] computeValueDistribution(int[] text) {
		double[] table1 = new double[getValueCount()];
		table1[getStartValue()] = 1.0;
		int state = getStartState();
		for (int c : text) {
			double[] table2 = new double[getValueCount()];
			state = transitionTarget(state, c);
			for (int value=0; value<getValueCount(); ++value) {
				if (table1[value]==0.0) continue;
				for (int emission=0; emission<getEmissionCount(); ++emission) {
					table2[performOperation(state, value, emission)] += table1[value]*emissionProbability(state, emission);
				}
			}
			table1 = table2;
		}
		return table1;
	}

	
}
